# %pylab inline
import matplotlib.pyplot as plt
import numpy as np
import os, sys, glob, shutil, json
# os用于操作文件和目录
# sys提供了与python解释器相关的功能
# glob用于文件通配符匹配
# shutil用于文件操作
# json用于处理JSON数据
from torch.optim.lr_scheduler import StepLR  # PyTorch中的学习率调度器，用于动态地调整模型的学习率。

# 设置CUDA设备
os.environ["CUDA_VISIBLE_DEVICES"] = '0'
import cv2
from PIL import Image
from tqdm import tqdm, tqdm_notebook  # 用于创建进度条以监视代码中循环的进展。
import torch

torch.manual_seed(0)  # 通过设置随机种子，可以实现重复性，便于调试和验证模型的稳定性。
torch.backends.cudnn.deterministic = False
torch.backends.cudnn.benchmark = True

import torchvision.models as models
import torchvision.transforms as transforms  # 用于图像预处理和数据增强的工具。例如，可以使用这些转换来对图像进行裁剪、缩放、标准化等操作。
import torchvision.datasets as datasets  # 用于加载常见图像数据集的工具
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torch.autograd import Variable
from torch.utils.data.dataset import Dataset
from torchvision.models import resnet18, ResNet18_Weights, resnet50, ResNet50_Weights


class SVHNDataset(Dataset):
    # 初始化数据集，接收图片路径列表img_path和标签列表img_label，以及可选的图像转换transform
    def __init__(self, img_path, img_label, transform=None):
        self.img_path = img_path
        self.img_label = img_label
        if transform is not None:
            self.transform = transform
        else:
            self.transform = None

    # 根据索引返回一张图片及其对应的标签。标签被填充到长度为5，不足的部分用10（表示空白字符）填充
    def __getitem__(self, index):
        img = Image.open(self.img_path[index]).convert('RGB')

        if self.transform is not None:
            img = self.transform(img)

        lbl = np.array(self.img_label[index], dtype=int)
        lbl = list(lbl) + (5 - len(lbl)) * [10]
        return img, torch.from_numpy(np.array(lbl[:5]))

    # 返回数据集中图片的数量
    def __len__(self):
        return len(self.img_path)


# 训练集和验证集
if __name__ == '__main__':
    train_path = glob.glob('../project/Data/mchar_train/*.png')
    train_path.sort()
    train_json = json.load(open('../project/Data/mchar_train.json'))
    train_label = [train_json[x]['label'] for x in train_json]

    train_loader = torch.utils.data.DataLoader(
        SVHNDataset(train_path, train_label,
                    transforms.Compose([
                        transforms.Resize((80, 160)),
                        transforms.RandomCrop((64, 128)),
                        transforms.ColorJitter(0.3, 0.3, 0.2),
                        transforms.RandomRotation(5),
                        transforms.ToTensor(),
                        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
                    ])),
        batch_size=40,
        shuffle=True,
        num_workers=3,
    )

    val_path = glob.glob('../project/Data/mchar_val/*.png')
    val_path.sort()
    val_json = json.load(open('../project/Data/mchar_val.json'))
    val_label = [val_json[x]['label'] for x in val_json]

    val_loader = torch.utils.data.DataLoader(
        SVHNDataset(val_path, val_label,
                    transforms.Compose([
                        transforms.Resize((80, 160)),
                        transforms.RandomCrop((64, 128)),
                        transforms.ColorJitter(0.3, 0.3, 0.2),
                        transforms.RandomRotation(5),
                        transforms.ToTensor(),
                        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
                    ])),
        batch_size=40,
        shuffle=False,
        num_workers=3,
    )


class SVHN_Model1(nn.Module):
    def __init__(self):
        super(SVHN_Model1, self).__init__()

        # resnet50
        model_conv = resnet50(weights=ResNet50_Weights.DEFAULT)
        model_conv.avgpool = nn.AdaptiveAvgPool2d(1)
        model_conv = nn.Sequential(*list(model_conv.children())[:-1])  # 去除最后一个fc layer
        self.cnn = model_conv

        self.hd_fc1 = nn.Linear(2048, 128)
        self.hd_fc2 = nn.Linear(2048, 128)
        self.hd_fc3 = nn.Linear(2048, 128)
        self.hd_fc4 = nn.Linear(2048, 128)
        self.hd_fc5 = nn.Linear(2048, 128)
        self.dropout_1 = nn.Dropout(0.25)
        self.dropout_2 = nn.Dropout(0.25)
        self.dropout_3 = nn.Dropout(0.25)
        self.dropout_4 = nn.Dropout(0.25)
        self.dropout_5 = nn.Dropout(0.25)
        self.fc1 = nn.Linear(128, 11)
        self.fc2 = nn.Linear(128, 11)
        self.fc3 = nn.Linear(128, 11)
        self.fc4 = nn.Linear(128, 11)
        self.fc5 = nn.Linear(128, 11)

    def forward(self, img):
        feat = self.cnn(img)
        feat = feat.view(feat.shape[0], -1)

        feat1 = self.hd_fc1(feat)
        feat2 = self.hd_fc2(feat)
        feat3 = self.hd_fc3(feat)
        feat4 = self.hd_fc4(feat)
        feat5 = self.hd_fc5(feat)
        feat1 = self.dropout_1(feat1)
        feat2 = self.dropout_2(feat2)
        feat3 = self.dropout_3(feat3)
        feat4 = self.dropout_4(feat4)
        feat5 = self.dropout_5(feat5)

        c1 = self.fc1(feat1)
        c2 = self.fc2(feat2)
        c3 = self.fc3(feat3)
        c4 = self.fc4(feat4)
        c5 = self.fc5(feat5)
        return c1, c2, c3, c4, c5



# 训练函数
def train(train_loader, model, criterion, optimizer, epoch):
    # 切换模型为训练模式
    model.train()
    train_loss = []

    for i, (input, target) in enumerate(train_loader):  # 每个批次大小是40
        input = input.cpu()
        target = target.cpu()

        target = target.long()
        c0, c1, c2, c3, c4 = model(input)
        loss = criterion(c0, target[:, 0]) + \
               criterion(c1, target[:, 1]) + \
               criterion(c2, target[:, 2]) + \
               criterion(c3, target[:, 3]) + \
               criterion(c4, target[:, 4])

        optimizer.zero_grad()  # 将优化器的梯度缓冲区清零，以准备计算新的梯度。
        loss.backward()  # 进行反向传播，计算梯度
        optimizer.step()  # 更新参数，通过梯度下降来最小化损失函数

        train_loss.append(loss.item())  # 将每一个批次的损失加入一个列表中
    return np.mean(train_loss)  # 返回每个批次的损失平均值


# 验证函数
def validate(val_loader, model, criterion):
    # 切换模型为预测模型
    model.eval()
    val_loss = []

    # 不记录模型梯度信息
    with torch.no_grad():
        for i, (input, target) in enumerate(val_loader):
            input = input.cpu()
            target = target.cpu()

            c0, c1, c2, c3, c4 = model(input)
            # 注意：将 target 转换为整数类型的张量
            target = target.long()

            loss = criterion(c0, target[:, 0]) + \
                   criterion(c1, target[:, 1]) + \
                   criterion(c2, target[:, 2]) + \
                   criterion(c3, target[:, 3]) + \
                   criterion(c4, target[:, 4])
            val_loss.append(loss.item())
    return np.mean(val_loss)


# 预测函数 数据增强
def predict(test_loader, model, tta=10):
    model.eval()
    test_pred_tta = None
    # TTA是一种在测试过程中对输入数据进行多次变换或扰动，并对每个变换后的输入进行预测，然后取多次预测的平均值以提高模型性能的技术。
    # TTA 次数
    for _ in range(tta):
        test_pred = []  # 存储每个测试样本的预测结果

        with torch.no_grad():
            for i, (input, target) in enumerate(test_loader):
                input = input.cpu()

                c0, c1, c2, c3, c4 = model(input)
                output = np.concatenate([
                    c0.data.numpy(),
                    c1.data.numpy(),
                    c2.data.numpy(),
                    c3.data.numpy(),
                    c4.data.numpy()], axis=1)

                test_pred.append(output)

        test_pred = np.vstack(test_pred)
        if test_pred_tta is None:
            test_pred_tta = test_pred
        else:
            test_pred_tta += test_pred

    return test_pred_tta


if __name__ == '__main__':
    # 定义初始学习率
    initial_lr = 0.001
    # 定义初始学习率
    model = SVHN_Model1()
    criterion = nn.CrossEntropyLoss()  # 交叉熵通常用于分类问题，适用于多类别分类任务。
    optimizer = torch.optim.Adam(model.parameters(), 0.001)
    # 创建学习率调度器
    scheduler = StepLR(optimizer, step_size=5, gamma=0.5)
    best_loss = 1000.0

    # 是否使用GPU
    use_cuda = torch.cuda.is_available()  # 检查CUDA是否可用

    if use_cuda:
        model = model.cuda()

    for epoch in range(10):
        train_loss = train(train_loader, model, criterion, optimizer, epoch)
        val_loss = validate(val_loader, model, criterion)

        # 在每个 epoch 结束时更新学习率
        scheduler.step()

        val_label = [''.join(map(str, x)) for x in val_loader.dataset.img_label]
        val_predict_label = predict(val_loader, model, 1)
        val_predict_label = np.vstack([
            val_predict_label[:, :11].argmax(1),
            val_predict_label[:, 11:22].argmax(1),
            val_predict_label[:, 22:33].argmax(1),
            val_predict_label[:, 33:44].argmax(1),
            val_predict_label[:, 44:55].argmax(1),
        ]).T
        val_label_pred = []
        for x in val_predict_label:
            val_label_pred.append(''.join(map(str, x[x != 10])))

        val_char_acc = np.mean(np.array(val_label_pred) == np.array(val_label))

        print('Epoch: {0}, Train loss: {1} \t Val loss: {2}'.format(epoch, train_loss, val_loss))
        print('Val Acc', val_char_acc)
        # 记录下验证集精度
        if val_loss < best_loss:
            best_loss = val_loss
            torch.save(model.state_dict(), './model.pt')

    test_path = glob.glob('../project/Data/mchar_test_a/*.png')
    test_path.sort()
    # test_json = json.load(open('../input/test_a.json'))
    test_label = [[1]] * len(test_path)
    # print(len(test_path), len(test_label))

    test_loader = torch.utils.data.DataLoader(
        SVHNDataset(test_path, test_label,
                    transforms.Compose([
                        transforms.Resize((68, 136)),
                        transforms.RandomCrop((64, 128)),
                        transforms.ColorJitter(0.3, 0.3, 0.2),
                        transforms.RandomRotation(5),
                        transforms.ToTensor(),
                        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
                    ])),
        batch_size=40,
        shuffle=False,
        num_workers=3,
    )

    # 加载保存的最优模型
    model.load_state_dict(torch.load('model.pt', weights_only=True))

    test_predict_label = predict(test_loader, model, 1)
    # print(test_predict_label.shape)

    test_label = [''.join(map(str, x)) for x in test_loader.dataset.img_label]
    test_predict_label = np.vstack([
        test_predict_label[:, :11].argmax(1),
        test_predict_label[:, 11:22].argmax(1),
        test_predict_label[:, 22:33].argmax(1),
        test_predict_label[:, 33:44].argmax(1),
        test_predict_label[:, 44:55].argmax(1),
    ]).T

    test_label_pred = []
    for x in test_predict_label:
        test_label_pred.append(''.join(map(str, x[x != 10])))

    import pandas as pd

    # 读取提交模板文件
    df_submit = pd.read_csv('../project/Data/mchar_sample_submit_A.csv')
    df_submit['file_code'] = test_label_pred
    df_submit.to_csv('../project/Data/submit.csv', index=None)
